#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于LangChain的天气查询智能体（支持国产模型）
用于查询城市天气并提供出行建议
"""

import json
import os
from typing import Optional
from langchain.agents import AgentType, initialize_agent
from langchain.tools import Tool
from langchain.llms.base import LLM
from langchain.schema import BaseMessage
from pydantic import Field
import requests


def get_weather(city: str) -> str:
    """
    模拟天气API，返回指定城市的天气信息
    
    Args:
        city: 城市名称
        
    Returns:
        包含天气信息的JSON字符串
    """
    weather_data = {
        "beijing": {
            "location": "Beijing",
            "temperature": {
                "current": 32,
                "low": 26,
                "high": 35
            },
            "rain_probability": 10,   # 百分比
            "humidity": 40  # 百分比
        },
        "shenzhen": {
            "location": "Shenzhen",
            "temperature": {
                "current": 28,
                "low": 24,
                "high": 31
            },
            "rain_probability": 90,   # 百分比
            "humidity": 85     # 百分比
        }
    }
    
    city_key = city.lower()
    if city_key in weather_data:
        return json.dumps(weather_data[city_key], ensure_ascii=False)
    return json.dumps({"error": "Weather Unavailable"}, ensure_ascii=False)


class DeepSeekLLM(LLM):
    """
    DeepSeek模型的LangChain封装
    """
    api_key: str = Field(default="")
    model_name: str = Field(default="deepseek-chat")
    temperature: float = Field(default=0.1)
    max_tokens: int = Field(default=500)
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        if not self.api_key:
            self.api_key = os.getenv("DEEPSEEK_API_KEY", "")
        if not self.api_key:
            raise ValueError("请设置DEEPSEEK_API_KEY环境变量")
    
    @property
    def _llm_type(self) -> str:
        return "deepseek"
    
    def _call(self, prompt: str, stop=None, **kwargs) -> str:
        """
        调用DeepSeek API
        """
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": self.temperature,
            "max_tokens": self.max_tokens
        }
        
        try:
            response = requests.post(
                "https://api.deepseek.com/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )
            response.raise_for_status()
            
            result = response.json()
            return result["choices"][0]["message"]["content"]
            
        except Exception as e:
            return f"API调用失败: {str(e)}"


def create_weather_agent_with_domestic_llm():
    """
    创建使用国产模型的天气查询智能体
    
    Returns:
        配置好的LangChain智能体
    """
    # 初始化国产LLM（以DeepSeek为例）
    llm = DeepSeekLLM(
        temperature=0.1,
        max_tokens=500
    )
    
    # 定义工具
    weather_tool = Tool(
        name="get_weather",
        description="获取指定城市的天气信息，包括温度、降雨概率和湿度。输入参数为城市名称（如：shenzhen, beijing）",
        func=get_weather
    )
    
    tools = [weather_tool]
    
    # 创建智能体
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        max_iterations=3,
        early_stopping_method="generate"
    )
    
    return agent


def main():
    """
    主函数：执行天气查询和出行建议任务
    """
    try:
        # 创建智能体
        agent = create_weather_agent_with_domestic_llm()
        
        # 定义任务
        task = "查找深圳的天气，然后根据天气情况用一句话告诉我出门要不要带伞"
        
        print("=== 智能体任务执行（国产模型版本）===")
        print(f"任务：{task}")
        print("\n=== 执行过程 ===")
        
        # 执行任务
        result = agent.run(task)
        
        print("\n=== 最终建议 ===")
        print(result)
        
    except Exception as e:
        print(f"执行出错：{str(e)}")
        print("请检查API密钥配置和网络连接")


if __name__ == "__main__":
    main()
